{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Image Recognition on MNIST using PyTorch Lightning\n",
    "\n",
    "Demonstrating the elements of machine learning:\n",
    "\n",
    "1) **E**xperience (Datasets and Dataloaders)<br>\n",
    "2) **T**ask (Classifier Model)<br>\n",
    "3) **P**erformance (Accuracy)<br>\n",
    "\n",
    "**Experience:** <br>\n",
    "We use MNIST dataset for this demo. MNIST is made of 28x28 images of handwritten digits, `0` to `9`. The train split has 60,000 images and the test split has 10,000 images. Images are all gray-scale.\n",
    "\n",
    "**Task:**<br>\n",
    "Our task is to classify the images into 10 classes. We use ResNet18 model from torchvision.models. The ResNet18 first convolutional layer (`conv1`) is modified to accept a single channel input. The number of classes is set to 10.\n",
    "\n",
    "**Performance:**<br>\n",
    "We use accuracy metric to evaluate the performance of our model on the test split. `torchmetrics.functional.accuracy`  calculates the accuracy.\n",
    "\n",
    "**[Pytorch Lightning](https://www.pytorchlightning.ai/):**<br>\n",
    "Our demo uses Pytorch Lightning to simplify the process of training and testing. Pytorch Lightning `Trainer` trains and evaluates our model. The default configurations are for a GPU-enabled system with 48 CPU cores. Please change the configurations if you have a different system.\n",
    "\n",
    "**[Weights and Biases](https://www.wandb.ai/):**<br>\n",
    "`wandb` is used by PyTorch Lightining Module to log train and evaluations results. Use `--no-wandb` to disable `wandb`.\n",
    "\n",
    "\n",
    "Let us install `pytorch-lightning` and `torchmetrics`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://pypi.ngc.nvidia.com\n",
      "Requirement already satisfied: pytorch-lightning in /home/rowel/anaconda3/lib/python3.7/site-packages (1.6.3)\n",
      "Collecting pytorch-lightning\n",
      "  Downloading pytorch_lightning-1.7.7-py3-none-any.whl (708 kB)\n",
      "\u001b[K     |████████████████████████████████| 708 kB 38.8 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: fsspec[http]!=2021.06.0,>=2021.05.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (2022.1.0)\n",
      "Requirement already satisfied: PyYAML>=5.4 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (5.4.1)\n",
      "Requirement already satisfied: tqdm>=4.57.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (4.62.3)\n",
      "Requirement already satisfied: torchmetrics>=0.7.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (0.8.2)\n",
      "Requirement already satisfied: numpy>=1.17.2 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (1.21.5)\n",
      "Collecting tensorboard>=2.9.1\n",
      "  Downloading tensorboard-2.10.1-py3-none-any.whl (5.9 MB)\n",
      "\u001b[K     |████████████████████████████████| 5.9 MB 93.5 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: packaging>=17.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (21.3)\n",
      "Requirement already satisfied: torch>=1.9.* in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (1.12.0)\n",
      "Requirement already satisfied: pyDeprecate>=0.3.1 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (0.3.2)\n",
      "Requirement already satisfied: typing-extensions>=4.0.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pytorch-lightning) (4.1.1)\n",
      "Requirement already satisfied: requests in /home/rowel/anaconda3/lib/python3.7/site-packages (from fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (2.27.1)\n",
      "Requirement already satisfied: aiohttp in /home/rowel/anaconda3/lib/python3.7/site-packages (from fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (3.7.4.post0)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /home/rowel/anaconda3/lib/python3.7/site-packages (from packaging>=17.0->pytorch-lightning) (3.0.4)\n",
      "Requirement already satisfied: absl-py>=0.4 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.9.1->pytorch-lightning) (0.12.0)\n",
      "Requirement already satisfied: google-auth<3,>=1.6.3 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.9.1->pytorch-lightning) (1.15.0)\n",
      "Collecting tensorboard-data-server<0.7.0,>=0.6.0\n",
      "  Downloading tensorboard_data_server-0.6.1-py3-none-manylinux2010_x86_64.whl (4.9 MB)\n",
      "\u001b[K     |████████████████████████████████| 4.9 MB 89.9 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: wheel>=0.26 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.9.1->pytorch-lightning) (0.37.1)\n",
      "Requirement already satisfied: grpcio>=1.24.3 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.9.1->pytorch-lightning) (1.32.0)\n",
      "Requirement already satisfied: setuptools>=41.0.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.9.1->pytorch-lightning) (59.5.0)\n",
      "Requirement already satisfied: markdown>=2.6.8 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.9.1->pytorch-lightning) (3.1.1)\n",
      "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.9.1->pytorch-lightning) (0.4.1)\n",
      "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.9.1->pytorch-lightning) (1.6.0.post3)\n",
      "Requirement already satisfied: protobuf<3.20,>=3.9.2 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.9.1->pytorch-lightning) (3.19.4)\n",
      "Requirement already satisfied: werkzeug>=1.0.1 in /home/rowel/anaconda3/lib/python3.7/site-packages (from tensorboard>=2.9.1->pytorch-lightning) (2.0.3)\n",
      "Requirement already satisfied: six in /home/rowel/anaconda3/lib/python3.7/site-packages (from absl-py>=0.4->tensorboard>=2.9.1->pytorch-lightning) (1.16.0)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /home/rowel/anaconda3/lib/python3.7/site-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning) (0.2.8)\n",
      "Requirement already satisfied: rsa<4.1,>=3.1.4 in /home/rowel/anaconda3/lib/python3.7/site-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning) (4.0)\n",
      "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning) (4.1.0)\n",
      "Requirement already satisfied: requests-oauthlib>=0.7.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.9.1->pytorch-lightning) (1.3.0)\n",
      "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /home/rowel/anaconda3/lib/python3.7/site-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning) (0.4.8)\n",
      "Requirement already satisfied: charset-normalizer~=2.0.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (2.0.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/rowel/anaconda3/lib/python3.7/site-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (2021.10.8)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/rowel/anaconda3/lib/python3.7/site-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (1.25.11)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /home/rowel/anaconda3/lib/python3.7/site-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (3.3)\n",
      "Requirement already satisfied: oauthlib>=3.0.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.9.1->pytorch-lightning) (3.1.0)\n",
      "Requirement already satisfied: chardet<5.0,>=2.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (4.0.0)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (1.6.3)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/rowel/anaconda3/lib/python3.7/site-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (5.1.0)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (21.4.0)\n",
      "Requirement already satisfied: async-timeout<4.0,>=3.0 in /home/rowel/anaconda3/lib/python3.7/site-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (3.0.1)\n",
      "Installing collected packages: tensorboard-data-server, tensorboard, pytorch-lightning\n",
      "  Attempting uninstall: tensorboard\n",
      "    Found existing installation: tensorboard 2.4.1\n",
      "    Uninstalling tensorboard-2.4.1:\n",
      "      Successfully uninstalled tensorboard-2.4.1\n",
      "  Attempting uninstall: pytorch-lightning\n",
      "    Found existing installation: pytorch-lightning 1.6.3\n",
      "    Uninstalling pytorch-lightning-1.6.3:\n",
      "      Successfully uninstalled pytorch-lightning-1.6.3\n",
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "tensorflow 2.4.1 requires numpy~=1.19.2, but you have numpy 1.21.5 which is incompatible.\n",
      "tensorflow 2.4.1 requires six~=1.15.0, but you have six 1.16.0 which is incompatible.\n",
      "tensorflow 2.4.1 requires typing-extensions~=3.7.4, but you have typing-extensions 4.1.1 which is incompatible.\n",
      "tensorflow 2.4.1 requires wrapt~=1.12.1, but you have wrapt 1.13.3 which is incompatible.\u001b[0m\n",
      "Successfully installed pytorch-lightning-1.7.7 tensorboard-2.10.1 tensorboard-data-server-0.6.1\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "Looking in indexes: https://pypi.org/simple, https://pypi.ngc.nvidia.com\n",
      "Requirement already satisfied: torchmetrics in /home/rowel/anaconda3/lib/python3.7/site-packages (0.8.2)\n",
      "Collecting torchmetrics\n",
      "  Downloading torchmetrics-0.10.0-py3-none-any.whl (529 kB)\n",
      "\u001b[K     |████████████████████████████████| 529 kB 32.4 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: torch>=1.3.1 in /home/rowel/anaconda3/lib/python3.7/site-packages (from torchmetrics) (1.12.0)\n",
      "Requirement already satisfied: numpy>=1.17.2 in /home/rowel/anaconda3/lib/python3.7/site-packages (from torchmetrics) (1.21.5)\n",
      "Requirement already satisfied: packaging in /home/rowel/anaconda3/lib/python3.7/site-packages (from torchmetrics) (21.3)\n",
      "Requirement already satisfied: typing-extensions in /home/rowel/anaconda3/lib/python3.7/site-packages (from torchmetrics) (4.1.1)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /home/rowel/anaconda3/lib/python3.7/site-packages (from packaging->torchmetrics) (3.0.4)\n",
      "Installing collected packages: torchmetrics\n",
      "  Attempting uninstall: torchmetrics\n",
      "    Found existing installation: torchmetrics 0.8.2\n",
      "    Uninstalling torchmetrics-0.8.2:\n",
      "      Successfully uninstalled torchmetrics-0.8.2\n",
      "Successfully installed torchmetrics-0.10.0\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install pytorch-lightning --upgrade\n",
    "%pip install torchmetrics --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import wandb\n",
    "from argparse import ArgumentParser\n",
    "from pytorch_lightning import LightningModule, Trainer, Callback\n",
    "from pytorch_lightning.loggers import WandbLogger\n",
    "from torchmetrics.functional import accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pytorch Lightning Module\n",
    "\n",
    "PyTorch Lightning Module has a PyTorch ResNet18 Model. It is a subclass of LightningModule. The model part is subclassed to support a single channel input. We replaced the input convolutional layer to support single channel inputs. The Lightning Module is also a container for the model, the optimizer, the loss function, the metrics, and the data loaders.\n",
    "\n",
    "`ResNet` class can be found [here](https://pytorch.org/vision/0.8/_modules/torchvision/models/resnet.html).\n",
    "\n",
    "By using PyTorch Lightning, we simplify the training and testing processes since we do not need to write boiler plate code blocks. These include automatic transfer to chosen device (i.e. `gpu` or `cpu`), model `eval` and `train` modes, and backpropagation routines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LitMNISTModel(LightningModule):\n",
    "    def __init__(self, num_classes=10, lr=0.001, batch_size=32):\n",
    "        super().__init__()\n",
    "        self.save_hyperparameters()\n",
    "        self.model = torchvision.models.resnet18(num_classes=num_classes)\n",
    "        self.model.conv1 = torch.nn.Conv2d(1, 64, kernel_size=7,\n",
    "                                           stride=2, padding=3, bias=False)\n",
    "        self.loss = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "\n",
    "    # this is called during fit()\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        y_hat = self.forward(x)\n",
    "        loss = self.loss(y_hat, y)\n",
    "        return {\"loss\": loss}\n",
    "\n",
    "    # calls to self.log() are recorded in wandb\n",
    "    def training_epoch_end(self, outputs):\n",
    "        avg_loss = torch.stack([x[\"loss\"] for x in outputs]).mean()\n",
    "        self.log(\"train_loss\", avg_loss, on_epoch=True)\n",
    "\n",
    "    # this is called at the end of an epoch\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        y_hat = self.forward(x)\n",
    "        loss = self.loss(y_hat, y)\n",
    "        acc = accuracy(y_hat, y) * 100.\n",
    "        # we use y_hat to display predictions during callback\n",
    "        return {\"y_hat\": y_hat, \"test_loss\": loss, \"test_acc\": acc}\n",
    "\n",
    "    # this is called at the end of all epochs\n",
    "    def test_epoch_end(self, outputs):\n",
    "        avg_loss = torch.stack([x[\"test_loss\"] for x in outputs]).mean()\n",
    "        avg_acc = torch.stack([x[\"test_acc\"] for x in outputs]).mean()\n",
    "        self.log(\"test_loss\", avg_loss, on_epoch=True, prog_bar=True)\n",
    "        self.log(\"test_acc\", avg_acc, on_epoch=True, prog_bar=True)\n",
    "\n",
    "    # validation is the same as test\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "       return self.test_step(batch, batch_idx)\n",
    "\n",
    "    def validation_epoch_end(self, outputs):\n",
    "        return self.test_epoch_end(outputs)\n",
    "\n",
    "    # we use Adam optimizer\n",
    "    def configure_optimizers(self):\n",
    "        return torch.optim.Adam(self.parameters(), lr=self.hparams.lr)\n",
    "    \n",
    "    # this is called after model instatiation to initiliaze the datasets and dataloaders\n",
    "    def setup(self, stage=None):\n",
    "        self.train_dataloader()\n",
    "        self.test_dataloader()\n",
    "\n",
    "    # build train and test dataloaders using MNIST dataset\n",
    "    # we use simple ToTensor transform\n",
    "    def train_dataloader(self):\n",
    "        return torch.utils.data.DataLoader(\n",
    "            torchvision.datasets.MNIST(\n",
    "                \"./data\", train=True, download=True, \n",
    "                transform=torchvision.transforms.ToTensor()\n",
    "            ),\n",
    "            batch_size=self.hparams.batch_size,\n",
    "            shuffle=True,\n",
    "            num_workers=48,\n",
    "            pin_memory=True,\n",
    "        )\n",
    "\n",
    "    def test_dataloader(self):\n",
    "        return torch.utils.data.DataLoader(\n",
    "            torchvision.datasets.MNIST(\n",
    "                \"./data\", train=False, download=True, \n",
    "                transform=torchvision.transforms.ToTensor()\n",
    "            ),\n",
    "            batch_size=self.hparams.batch_size,\n",
    "            shuffle=False,\n",
    "            num_workers=48,\n",
    "            pin_memory=True,\n",
    "        )\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return self.test_dataloader()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PyTorch Lightning Callback\n",
    "\n",
    "We can instantiate a callback object to perform certain tasks during training. In this case, we log sample images, ground truth labels, and predicted labels from the test dataset.\n",
    "\n",
    "We can also `ModelCheckpoint` callback to save the model after each epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WandbCallback(Callback):\n",
    "\n",
    "    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):\n",
    "        # process first 10 images of the first batch\n",
    "        if batch_idx == 0:\n",
    "            n = 10\n",
    "            x, y = batch\n",
    "            outputs = outputs[\"y_hat\"]\n",
    "            outputs = torch.argmax(outputs, dim=1)\n",
    "            # log image, ground truth and prediction on wandb table\n",
    "            columns = ['image', 'ground truth', 'prediction']\n",
    "            data = [[wandb.Image(x_i), y_i, y_pred] for x_i, y_i, y_pred in list(\n",
    "                zip(x[:n], y[:n], outputs[:n]))]\n",
    "            wandb_logger.log_table(\n",
    "                key='ResNet18 on MNIST Predictions',\n",
    "                columns=columns,\n",
    "                data=data)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Program Arguments\n",
    "\n",
    "When running on command line, we can pass arguments to the program. For the jupyter notebook, we can pass arguments using the `%run` magic command.\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_args():\n",
    "    parser = ArgumentParser(description=\"PyTorch Lightning MNIST Example\")\n",
    "    parser.add_argument(\"--max-epochs\", type=int, default=5, help=\"num epochs\")\n",
    "    parser.add_argument(\"--batch-size\", type=int, default=32, help=\"batch size\")\n",
    "    parser.add_argument(\"--lr\", type=float, default=0.001, help=\"learning rate\")\n",
    "\n",
    "    parser.add_argument(\"--num-classes\", type=int, default=10, help=\"num classes\")\n",
    "\n",
    "    parser.add_argument(\"--devices\", default=1)\n",
    "    parser.add_argument(\"--accelerator\", default='gpu')\n",
    "    parser.add_argument(\"--num-workers\", type=int, default=48, help=\"num workers\")\n",
    "    \n",
    "    parser.add_argument(\"--no-wandb\", default=False, action='store_true')\n",
    "    args = parser.parse_args(\"\")\n",
    "    return args"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training and Evaluation using `Trainer`\n",
    "\n",
    "Get command line arguments. Instatiate a Pytorch Lightning Model. Train the model. Evaluate the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LitMNISTModel(\n",
      "  (model): ResNet(\n",
      "    (conv1): Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
      "    (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (relu): ReLU(inplace=True)\n",
      "    (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
      "    (layer1): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (layer2): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (downsample): Sequential(\n",
      "          (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (layer3): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (downsample): Sequential(\n",
      "          (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (layer4): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (downsample): Sequential(\n",
      "          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
      "    (fc): Linear(in_features=512, out_features=10, bias=True)\n",
      "  )\n",
      "  (loss): CrossEntropyLoss()\n",
      ")\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mrowel\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n",
      "Exception in thread StatsThr:\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/rowel/anaconda3/envs/voice/lib/python3.9/threading.py\", line 973, in _bootstrap_inner\n",
      "    self.run()\n",
      "  File \"/home/rowel/anaconda3/envs/voice/lib/python3.9/threading.py\", line 910, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/rowel/anaconda3/envs/voice/lib/python3.9/site-packages/wandb/sdk/internal/stats.py\", line 137, in _thread_body\n",
      "    stats = self.stats()\n",
      "  File \"/home/rowel/anaconda3/envs/voice/lib/python3.9/site-packages/wandb/sdk/internal/stats.py\", line 183, in stats\n",
      "    handle = pynvml.nvmlDeviceGetHandleByIndex(i)\n",
      "  File \"/home/rowel/anaconda3/envs/voice/lib/python3.9/site-packages/wandb/vendor/pynvml/pynvml.py\", line 2232, in nvmlDeviceGetHandleByIndex\n",
      "    _nvmlCheckReturn(ret)\n",
      "  File \"/home/rowel/anaconda3/envs/voice/lib/python3.9/site-packages/wandb/vendor/pynvml/pynvml.py\", line 1015, in _nvmlCheckReturn\n",
      "    raise NVMLError(ret)\n",
      "wandb.vendor.pynvml.pynvml.NVMLError_Unknown: Unknown Error\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.13.4"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/home/rowel/github/roatienza/Deep-Learning-Experiments/versions/2022/supervised/python/wandb/run-20221011_084931-3pxvztmt</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href=\"https://wandb.ai/rowel/pl-mnist/runs/3pxvztmt\" target=\"_blank\">cool-grass-7</a></strong> to <a href=\"https://wandb.ai/rowel/pl-mnist\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://wandb.me/run\" target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "MisconfigurationException",
     "evalue": "GPUAccelerator can not run on your system since the accelerator is not available. The following accelerator(s) is available and can be passed into `accelerator` argument of `Trainer`: ['cpu'].",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mMisconfigurationException\u001b[0m                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [5], line 13\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[38;5;66;03m# wandb is a great way to debug and visualize this model\u001b[39;00m\n\u001b[1;32m     11\u001b[0m wandb_logger \u001b[38;5;241m=\u001b[39m WandbLogger(project\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpl-mnist\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m---> 13\u001b[0m trainer \u001b[38;5;241m=\u001b[39m \u001b[43mTrainer\u001b[49m\u001b[43m(\u001b[49m\u001b[43maccelerator\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43maccelerator\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     14\u001b[0m \u001b[43m                  \u001b[49m\u001b[43mdevices\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdevices\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     15\u001b[0m \u001b[43m                  \u001b[49m\u001b[43mmax_epochs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmax_epochs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     16\u001b[0m \u001b[43m                  \u001b[49m\u001b[43mlogger\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwandb_logger\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mnot\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mno_wandb\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     17\u001b[0m \u001b[43m                  \u001b[49m\u001b[43mcallbacks\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m[\u001b[49m\u001b[43mWandbCallback\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mnot\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mno_wandb\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     18\u001b[0m trainer\u001b[38;5;241m.\u001b[39mfit(model)\n\u001b[1;32m     19\u001b[0m trainer\u001b[38;5;241m.\u001b[39mtest(model)\n",
      "File \u001b[0;32m~/anaconda3/envs/voice/lib/python3.9/site-packages/pytorch_lightning/utilities/argparse.py:339\u001b[0m, in \u001b[0;36m_defaults_from_env_vars.<locals>.insert_env_defaults\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    336\u001b[0m kwargs \u001b[39m=\u001b[39m \u001b[39mdict\u001b[39m(\u001b[39mlist\u001b[39m(env_variables\u001b[39m.\u001b[39mitems()) \u001b[39m+\u001b[39m \u001b[39mlist\u001b[39m(kwargs\u001b[39m.\u001b[39mitems()))\n\u001b[1;32m    338\u001b[0m \u001b[39m# all args were already moved to kwargs\u001b[39;00m\n\u001b[0;32m--> 339\u001b[0m \u001b[39mreturn\u001b[39;00m fn(\u001b[39mself\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
      "File \u001b[0;32m~/anaconda3/envs/voice/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:485\u001b[0m, in \u001b[0;36mTrainer.__init__\u001b[0;34m(self, logger, checkpoint_callback, enable_checkpointing, callbacks, default_root_dir, gradient_clip_val, gradient_clip_algorithm, process_position, num_nodes, num_processes, devices, gpus, auto_select_gpus, tpu_cores, ipus, log_gpu_memory, progress_bar_refresh_rate, enable_progress_bar, overfit_batches, track_grad_norm, check_val_every_n_epoch, fast_dev_run, accumulate_grad_batches, max_epochs, min_epochs, max_steps, min_steps, max_time, limit_train_batches, limit_val_batches, limit_test_batches, limit_predict_batches, val_check_interval, flush_logs_every_n_steps, log_every_n_steps, accelerator, strategy, sync_batchnorm, precision, enable_model_summary, weights_summary, weights_save_path, num_sanity_val_steps, resume_from_checkpoint, profiler, benchmark, deterministic, reload_dataloaders_every_n_epochs, auto_lr_find, replace_sampler_ddp, detect_anomaly, auto_scale_batch_size, prepare_data_per_node, plugins, amp_backend, amp_level, move_metrics_to_cpu, multiple_trainloader_mode, stochastic_weight_avg, terminate_on_nan)\u001b[0m\n\u001b[1;32m    482\u001b[0m \u001b[39m# init connectors\u001b[39;00m\n\u001b[1;32m    483\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_data_connector \u001b[39m=\u001b[39m DataConnector(\u001b[39mself\u001b[39m, multiple_trainloader_mode)\n\u001b[0;32m--> 485\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_accelerator_connector \u001b[39m=\u001b[39m AcceleratorConnector(\n\u001b[1;32m    486\u001b[0m     num_processes\u001b[39m=\u001b[39;49mnum_processes,\n\u001b[1;32m    487\u001b[0m     devices\u001b[39m=\u001b[39;49mdevices,\n\u001b[1;32m    488\u001b[0m     tpu_cores\u001b[39m=\u001b[39;49mtpu_cores,\n\u001b[1;32m    489\u001b[0m     ipus\u001b[39m=\u001b[39;49mipus,\n\u001b[1;32m    490\u001b[0m     accelerator\u001b[39m=\u001b[39;49maccelerator,\n\u001b[1;32m    491\u001b[0m     strategy\u001b[39m=\u001b[39;49mstrategy,\n\u001b[1;32m    492\u001b[0m     gpus\u001b[39m=\u001b[39;49mgpus,\n\u001b[1;32m    493\u001b[0m     num_nodes\u001b[39m=\u001b[39;49mnum_nodes,\n\u001b[1;32m    494\u001b[0m     sync_batchnorm\u001b[39m=\u001b[39;49msync_batchnorm,\n\u001b[1;32m    495\u001b[0m     benchmark\u001b[39m=\u001b[39;49mbenchmark,\n\u001b[1;32m    496\u001b[0m     replace_sampler_ddp\u001b[39m=\u001b[39;49mreplace_sampler_ddp,\n\u001b[1;32m    497\u001b[0m     deterministic\u001b[39m=\u001b[39;49mdeterministic,\n\u001b[1;32m    498\u001b[0m     auto_select_gpus\u001b[39m=\u001b[39;49mauto_select_gpus,\n\u001b[1;32m    499\u001b[0m     precision\u001b[39m=\u001b[39;49mprecision,\n\u001b[1;32m    500\u001b[0m     amp_type\u001b[39m=\u001b[39;49mamp_backend,\n\u001b[1;32m    501\u001b[0m     amp_level\u001b[39m=\u001b[39;49mamp_level,\n\u001b[1;32m    502\u001b[0m     plugins\u001b[39m=\u001b[39;49mplugins,\n\u001b[1;32m    503\u001b[0m )\n\u001b[1;32m    504\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_logger_connector \u001b[39m=\u001b[39m LoggerConnector(\u001b[39mself\u001b[39m, log_gpu_memory)\n\u001b[1;32m    505\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_callback_connector \u001b[39m=\u001b[39m CallbackConnector(\u001b[39mself\u001b[39m)\n",
      "File \u001b[0;32m~/anaconda3/envs/voice/lib/python3.9/site-packages/pytorch_lightning/trainer/connectors/accelerator_connector.py:201\u001b[0m, in \u001b[0;36mAcceleratorConnector.__init__\u001b[0;34m(self, devices, num_nodes, accelerator, strategy, plugins, precision, amp_type, amp_level, sync_batchnorm, benchmark, replace_sampler_ddp, deterministic, auto_select_gpus, num_processes, tpu_cores, ipus, gpus)\u001b[0m\n\u001b[1;32m    199\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_accelerator_flag \u001b[39m==\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mauto\u001b[39m\u001b[39m\"\u001b[39m \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_accelerator_flag \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    200\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_accelerator_flag \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_choose_accelerator()\n\u001b[0;32m--> 201\u001b[0m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_set_parallel_devices_and_init_accelerator()\n\u001b[1;32m    203\u001b[0m \u001b[39m# 3. Instantiate ClusterEnvironment\u001b[39;00m\n\u001b[1;32m    204\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mcluster_environment: ClusterEnvironment \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_choose_and_init_cluster_environment()\n",
      "File \u001b[0;32m~/anaconda3/envs/voice/lib/python3.9/site-packages/pytorch_lightning/trainer/connectors/accelerator_connector.py:519\u001b[0m, in \u001b[0;36mAcceleratorConnector._set_parallel_devices_and_init_accelerator\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    515\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39maccelerator\u001b[39m.\u001b[39mis_available():\n\u001b[1;32m    516\u001b[0m     available_accelerator \u001b[39m=\u001b[39m [\n\u001b[1;32m    517\u001b[0m         acc_str \u001b[39mfor\u001b[39;00m acc_str \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_accelerator_types \u001b[39mif\u001b[39;00m AcceleratorRegistry\u001b[39m.\u001b[39mget(acc_str)\u001b[39m.\u001b[39mis_available()\n\u001b[1;32m    518\u001b[0m     ]\n\u001b[0;32m--> 519\u001b[0m     \u001b[39mraise\u001b[39;00m MisconfigurationException(\n\u001b[1;32m    520\u001b[0m         \u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m{\u001b[39;00m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39maccelerator\u001b[39m.\u001b[39m\u001b[39m__class__\u001b[39m\u001b[39m.\u001b[39m\u001b[39m__qualname__\u001b[39m\u001b[39m}\u001b[39;00m\u001b[39m can not run on your system\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m    521\u001b[0m         \u001b[39m\"\u001b[39m\u001b[39m since the accelerator is not available. The following accelerator(s)\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m    522\u001b[0m         \u001b[39m\"\u001b[39m\u001b[39m is available and can be passed into `accelerator` argument of\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m    523\u001b[0m         \u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m `Trainer`: \u001b[39m\u001b[39m{\u001b[39;00mavailable_accelerator\u001b[39m}\u001b[39;00m\u001b[39m.\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m    524\u001b[0m     )\n\u001b[1;32m    526\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_set_devices_flag_if_auto_passed()\n\u001b[1;32m    528\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_gpus \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_devices_flag \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_gpus \u001b[39melse\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_gpus\n",
      "\u001b[0;31mMisconfigurationException\u001b[0m: GPUAccelerator can not run on your system since the accelerator is not available. The following accelerator(s) is available and can be passed into `accelerator` argument of `Trainer`: ['cpu']."
     ]
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    args = get_args()\n",
    "    model = LitMNISTModel(num_classes=args.num_classes,\n",
    "                          lr=args.lr, batch_size=args.batch_size)\n",
    "    model.setup()\n",
    "\n",
    "    # printing the model is useful for debugging\n",
    "    print(model)\n",
    "\n",
    "    # wandb is a great way to debug and visualize this model\n",
    "    wandb_logger = WandbLogger(project=\"pl-mnist\")\n",
    "    \n",
    "    trainer = Trainer(accelerator=args.accelerator,\n",
    "                      devices=args.devices,\n",
    "                      max_epochs=args.max_epochs,\n",
    "                      logger=wandb_logger if not args.no_wandb else None,\n",
    "                      callbacks=[WandbCallback() if not args.no_wandb else None])\n",
    "    trainer.fit(model)\n",
    "    trainer.test(model)\n",
    "\n",
    "    wandb.finish()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.12 ('voice')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "975196f818cef4d8418ca9caaf131be767e393cdee268c01707e7bb8d73c6de2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
